﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.IO;

namespace Kokomo.Mirror.PeFormat
{
	public class ImportDirectory
	{
		public static ImportDirectory LoadFromModule(PeModule module)
		{
			if (module == null) throw new ArgumentNullException("module");

			DataDirectoryEntry importDirectoryEntry = module.ImportDirectoryEntry;
			if (importDirectoryEntry == null)
				throw new ArgumentException("The module does not contain an import directory", "module");
			DataDirectoryEntry delayLoadImportDirectoryEntry = module.DelayLoadImportDirectoryEntry;

			StructSerializer serializer = new StructSerializer();
			ImportDirectory importDirectory = new ImportDirectory();

			using (Stream stream = module.GetStream())
			{
				stream.Position = module.RvaToFileOffset(importDirectoryEntry.VirtualAddress).Value;
				LoadImportDescriptors(module, stream, importDirectory);

				if (delayLoadImportDirectoryEntry != null){
				stream.Position = module.RvaToFileOffset(delayLoadImportDirectoryEntry.VirtualAddress).Value;
					LoadDelayLoadImportDescriptors(module, stream, importDirectory);
			}}

			return importDirectory;
		}
		private static void LoadImportDescriptors(
			PeModule module,
			Stream stream,
			ImportDirectory importDirectory
			)
		{
			StructSerializer serializer = new StructSerializer();
			ModuleImportDescriptor[] moduleImportDescriptors = serializer.ReadArray<ModuleImportDescriptor>(null, stream);

			BinaryReader reader = new BinaryReader(stream);
			for (int iModule = 0; iModule < moduleImportDescriptors.Length; iModule++)
			{
				ModuleImportDescriptor moduleImportDescriptor = moduleImportDescriptors[iModule];

				stream.Position = module.RvaToFileOffset(moduleImportDescriptor.Name).Value;
				string moduleName = reader.ReadNullTerminatedString();
				ImportedModule importedModule = new ImportedModule(moduleName, moduleImportDescriptor.TimeDateStamp);
				importDirectory.ImportedModules.Add(importedModule);

				if (module.Is32Bit)
				{
					LoadImportTable(
						stream,
						module,
						moduleImportDescriptor.FirstThunk,
						moduleImportDescriptor.OriginalFirstThunk,
						importedModule.Functions
						);
				}
				else if (module.Is64Bit)
				{
					long[] thunks = serializer.ReadArray<long>(null, stream);
				}
			}
		}
		private static void LoadDelayLoadImportDescriptors(
			PeModule module,
			Stream stream,
			ImportDirectory importDirectory
			)
		{
			StructSerializer serializer = new StructSerializer();

			ImgDelayDescr[] delayDescriptors = serializer.ReadArray<ImgDelayDescr>(null, stream);

			BinaryReader reader = new BinaryReader(stream);

			for (int i = 0; i < delayDescriptors.Length; i++)
			{
				ImgDelayDescr delayDescriptor = delayDescriptors[i];

				stream.Position = module.RvaToFileOffset(delayDescriptor.rvaDLLName).Value;
				string moduleName = reader.ReadNullTerminatedString();
				ImportedModule importedModule = new ImportedModule(moduleName, (int)delayDescriptor.dwTimeStamp, true);
				importDirectory.ImportedModules.Add(importedModule);

				uint iatRva = delayDescriptor.rvaIat;
				uint intRva = delayDescriptor.rvaInt;
				if (!delayDescriptor.grAttrs.HasFlag(DelayloadAttributes.Rva))
				{
					iatRva -= (uint)module.OptionalHeader.ImageBase;
					intRva -= (uint)module.OptionalHeader.ImageBase;
				}

				LoadImportTable(
					stream,
					module,
					iatRva,
					intRva,
					importedModule.Functions
					);
			}
		}
		private static IList<ImportedFunction> LoadImportTable(
			Stream stream,
			PeModule module,
			uint iatRva,
			uint intRva,
			IList<ImportedFunction> functions
			)
		{
			uint firstThunkRva = iatRva;
			uint originalFirstThunkRva = intRva;

			StructSerializer serializer = new StructSerializer();

			int[] importDescriptorRvas = null;
			if (originalFirstThunkRva != 0)
			{
				stream.Position = module.RvaToFileOffset(originalFirstThunkRva).Value;
				importDescriptorRvas = serializer.ReadArray<int>(null, stream);
			}
			else if (firstThunkRva != 0)
			{
				stream.Position = module.RvaToFileOffset(firstThunkRva).Value;
				importDescriptorRvas = serializer.ReadArray<int>(null, stream);
			}

			int functionCount;
			if (importDescriptorRvas != null) functionCount = importDescriptorRvas.Length;
			else functionCount = 0;

			for (int iFunction = 0; iFunction < importDescriptorRvas.Length; iFunction++)
			{
				//if (
				//    (originalThunkAddresses != null)
				//    && (iFunction < originalThunkAddresses.Length)
				//    )
				//{
				//    originalThunkAddress = originalThunkAddresses[iFunction];
				//}

				uint thunkRva = firstThunkRva + (uint)iFunction * 4;

				int descriptorRva = importDescriptorRvas[iFunction];

				string name = null;
				int ordinal = -1;
				if (descriptorRva < 0)
				{
					ordinal = (int)(descriptorRva & 0x7FFFFFFF);
				}
				else
				{
					uint? fileOffset = module.RvaToFileOffset((uint)descriptorRva);
					if (fileOffset.HasValue)
					{
						stream.Position = fileOffset.Value;
						ImportByNameData importByName = serializer.ReadStruct<ImportByNameData>(stream);
						ordinal = importByName.Hint;
						name = importByName.Name;
					}
				}
				ImportedFunction function = new ImportedFunction(
					ordinal,
					name,
					//originalThunkAddress,
					thunkRva
					);

				functions.Add(function);
			}

			return functions;
		}

		public ImportDirectory()
		{
			this.ImportedModules = new List<ImportedModule>();
		}

		public IList<ImportedModule> ImportedModules { get; private set; }
	}
}
